from email.mime import image
from time import sleep
import numpy as np
from configure import configure
from Evaluation import DemoEvaluation, Evaluation, DemoTest
from CMPSO import CMPSO
from CMDE import CMDE
import os
import torch
from utils import dict_show, Plot
from CE import CMEA
import torch

config = vars(configure())

os.environ['CUDA_VISBLE_DEVICES'] = '{}'.format(config['gpu_id'])
device = 'cuda' if torch.cuda.is_available() else 'cpu'

param = {
    'dataset': config['dataset'],
    'lr': config['lr'],
    'batch_size': config['batch_size'],
    'epochs': config['epochs'],
    'model_arch': config['model_arch'],
    'device': device,
    'eps': config['eps'],
}

evaluationer = Evaluation(param=param)
# evaluationer = DemoEvaluation(param=param)

config['evaluationer'] = evaluationer
config['DsDim'] = evaluationer.get_decision_sapce_dim()
# config['DsDim'] = evaluationer.demoer.get_param()['DsDim']
dict_show(config)



"""
Init, which include the following setp:
1) randomly init swarms, including 0~OsDim-1, the i-th swarm includes DsDim particles.
2) eval the function value of th eparticles from different swarms
3) update pBest, gBest
4) update Arachieve
"""
# cmpsoer = CMPSO(config=config)
# cmpsoer = CMDE(config=config)
cmeaer = CMEA(config=config)

gen = 0

ploter = Plot()

# ## mainly iteration algorithm body
# while evaluationer.fes < config['FES']:

#     ## Generations
#     gen += 1
#     print('Generation: {}'.format(gen))
#     print('Current fes: {}'.format(evaluationer.fes))

#     ## inertia weight \w to be decrease linearly from 0.9 to 0.4 
#     config['inertiaw'] = 0.9 - (0.9 - 0.4) * (evaluationer.fes / config['FES'])
#     print('Current inertiaw in main: {:.2f}'.format(config['inertiaw']))
    
#     ## update the particles in different swarms
#     cmpsoer.update_swarms()

#     ## update the archive  
#     cmpsoer.update_archive()

#     print('Current archive size : {}'.format(cmpsoer.archive['size']))

#     ploter.plot_archive(cmpsoer.archive)

#     print()

flag = True
while evaluationer.fes < config['FES']:

    gen += 1
    print('Generation: {}'.format(gen))
    print('Current fes: {}'.format(evaluationer.fes))


    if evaluationer.fes < 1000:

        cmeaer.init_solution()
    else:

        if flag:
            flag = False
            cmeaer.init()

        cmeaer.update_populations()

        cmeaer.update_archive()
        print('Current archive size : {}'.format(cmeaer.archive['size']))

        ploter.plot_archive(cmeaer.archive)
# images = []
# images.append(evaluationer.get_figure(isClean=True))

# for individual in cmeaer.population:
#     images.append(evaluationer.get_figure(individual.x))

# images = torch.cat(images, dim=0)

# ploter.plot_image(images, name='image0')

